#!/usr/bin/python
from rap_core import *
from optparse import OptionParser
import os.path
import sys
import re

# Description:
#   Selects levels from a RAP file (adf04, olg, TCCDW.DAT, or NIST) according to
#   specified filters.  The filters can choose on any combination of 
#   configuration, J, parity, term, and energy.  The resulting list of level
#   indices can be read in to other RTCs to focus the application on that set.
#
#   This function is also useful for getting summary information from a file
#   using the -i option and looking at a list of energies that have been
#   filtered using the -e option.
#
#   The filters take the form:
#     configuration -  c1 c3     (see -i option)
#     term          -  2P 4S
#     2J            -  3  6
#     parity        -  even odd
#     energy range  -  [0.0:10.0]  [4,6]
#
# Note: 
#   when several filters of the same type are included in the filter list, the 
#   effects are OR'ed together.  Filters of different types are AND'ed.  An 
#   exception exists for the energy range; only the last energy range specified
#   will be used.
#
# Options:
#   -h, --help   - print help screen
#   -i, --info   - print information about adf04 file including
#                  configuration identifiers
#   -e, --energy - print level energies to screen
#   -l FILE, --levlist=FILE - input level list (for further filtering)d
#   -o FILE, --output=FILE  - output file name
#   -f FILE, --filter=FILE  - input file containing filter list 
#   --eV         - ouput energies in electron-volts
#

def report_info(ion,levlist,filename=None,useev=False):
  """ Output general information """
  (theight,twidth)=rlf_termsize()         # get height and width of screen

  # format configuration output (left column)
  lout=[]
  lout.append('%4i'%len(levlist)+' levels')
  lout.append('')
  lout.append('Configurations: ')
  for i in range(len(confs)):
    c=confs[i]
    lout.append('c%-2i'%(i+1)+'  '+str(c))

  # format energy, J, and term information (right column)
  rout=[]
  rout.append('')
  rout.append('')
  if ion.core_energy():
    if useev:
      rout.append('Core Energy (eV):')
      rout.append('%18.6f'%(RYDBERG*ion.core_energy()))
    else:
      rout.append('Core Energy (Ry):')
      rout.append('%18.6f'%ion.core_energy())
    rout.append('')
  if ion.ionization_threshold():
    if useev:
      rout.append('Ionization Threshold (eV):')
      rout.append('%12.6f'%(RYDBERG*ion.ionization_threshold()))
    else:
      rout.append('Ionization Threshold (Ry):')
      rout.append('%12.6f'%ion.ionization_threshold())
    rout.append('')
  if useev:
    const=RYDBERG
    rout.append('Energy Range (eV):')
  else:
    const=1.0
    rout.append('Energy Range (Ry):')
  semin=('%12.6f'%(const*emin)).strip()
  semax=('%12.6f'%(const*emax)).strip()
  rout.append('  ['+semin+','+semax+']')
  if cup == 'IC':
    rout.append('')
    rout.append('2J Range:')
    sjmin=('%2i'%round(2.0*jmin)).strip()
    sjmax=('%2i'%round(2.0*jmax)).strip()
    rout.append('  ('+sjmin+','+sjmax+')')
  rout.append('')
  rout.append('Terms:')
  while len(terms) > 0:
    line=''
    for i in range(5):
      if len(terms) == 0: break
      line=line+'%4s'%terms.pop(0)
    rout.append(line)

  outlines=[]           # store output lines
  outbuff=5             # separation between columns
  leftlen=0
  for line in lout: leftlen=max(leftlen,len(line))     # width of left column
  rightlen=0
  for line in rout: rightlen=max(rightlen,len(line))   # width of right column
  lform='%-'+str(leftlen)+'s'
  rform='%-'+str(rightlen)+'s'
  bform='%-'+str(outbuff)+'s'
  if (leftlen+rightlen+outbuff) < twidth:    # output as two-columns
    buffstr=bform%' '
    if len(lout) >= len(rout):
      while len(rout) > 0:
        outlines.append(lform%(lout.pop(0))+buffstr+rform%(rout.pop(0)))
      while len(lout) > 0: outlines.append(lform%(lout.pop(0)))
    else:
      leftstr=lform%' '
      while len(lout) > 0: 
        outlines.append(lform%lout.pop(0)+buffstr+rform%rout.pop(0))
      while len(rout) > 0: outlines.append(leftstr+buffstr+rform%rout.pop(0))
  else:                                      # output as one-column
    for line in lout: outlines.append(line)
    outlines.append('')
    for line in rout: outlines.append(line)
  if filename == None:
    for line in outlines: print line
  else:
    fout=open(filename,'w')
    for line in outlines: fout.write(line+'\n')
    fout.close()
  return True

def report_energies(filename=None,useev=False):
  """ Output level list """

  ibuff=10          # column buffer
  clen=0
  for c in confs: clen=max(clen,len(str(c)))   # get max length of configuration

  # generate format strings
  iform='%-7i'
  cwidth=clen+12
  cform='%-'+str(cwidth)+'s'
  eform='%12.6f'
  bform='%'+str(ibuff)+'s'
  bstr=bform%' '
  totlen=len(iform%1+cform%' '+eform%1.1)
  
  # generate output arrays
  tout=[]
  tout.append('%4i'%len(levlist)+' levels')
  tout.append('')
  hout='idx'.ljust(len(iform%1))
  hout=hout+'level'.ljust(len(cform%' '))
  title='energy (Ry)'
  if useev: title='energy (eV)'
  hout=hout+title.rjust(len(eform%1.1))
  out=[]
  for ilev in levlist: 
    istr=iform%ilev
    cstr=cform%ion.label(ilev)
    e=ion.energy(ilev)
    if useev: e*=RYDBERG
    estr=eform%e
    out.append(istr+cstr+estr)

  outlines=tout
  (theight,twidth)=rlf_termsize()                  # height and width of screen
  nlev=len(levlist)
  if nlev > theight and 2*totlen+ibuff < twidth:   # check if two-column output
    nlev2=nlev/2
    ileft=nlev%2
    iadd=nlev2+ileft
    outlines.append(hout+bstr+hout)
    for i in range(nlev2): outlines.append(out[i]+bstr+out[i+iadd])
    if ileft == 1: outlines.append(out[nlev2])
  else:
    outlines.append(hout)
    for line in out: outlines.append(line)
  if filename == None:
    for line in outlines: print line
  else:
    fout=open(filename,'w')
    for line in outlines: fout.write(line+'\n')
    fout.close()
  return True

def apply_filter(args,levlist=None):
  # initialize filter variables
  cfilt=[]
  tfilt=[]
  par=None
  j2filt=[]
  emin=None
  emax=None

  # sort filter arguments
  for a in args:
    if not isinstance(a,str): continue
    if a[0].upper() == 'C':              # configuration
      try:
        i=int(a[1:])
        if i < 1 or i > len(confs):
          out=[]
          return out
        else:
          cfilt.append(confs[i-1])
      except:
        continue
    elif a in terms:                     # term
      tfilt.append(a)
      continue
    elif a.upper() == 'EVEN':            # even parity
      par=0
      continue
    elif a.upper() == 'ODD':             # odd parity
      par=1
      continue
    else:
      tmp=re.findall(r'\[\s*(\d+\.*\d*)\s*[,:]\s*(\d+\.*\d*)\s*\]',a)
      if len(tmp) > 0:                     # energy range
        temin=float(tmp[0][0])
        temax=float(tmp[0][1])
        if temin <= temax: emin,emax=temin,temax
      else:
        try:                                 # 2J
          j2=int(a)
        except:
          out=[]
          return out
        else:
          j2filt.append(j2)
          continue

  # apply filters to input level list
  try:
    out=levlist[:]
  except:
    return False
  if len(j2filt) > 0: out=ion.j2filter(j2filt,subset=out)
  if out == False: return False
  if len(cfilt) > 0: out=ion.cfilter(cfilt,subset=out)
  if out == False: return False
  if len(tfilt) > 0: out=ion.tfilter(tfilt,subset=out)
  if out == False: return False
  if par != None: out=ion.pfilter(par,subset=out)
  if out == False: return False
  if emin != None and emax != None: 
    if useev:
      out=ion.efilter(emin/RYDBERG,emax/RYDBERG,subset=out)
    else:
      out=ion.efilter(emin,emax,subset=out)
  if out == False: return False

  return out

def write_levlist(levlist,filename=None):
  # make copy of level list
  try:
    lprint=levlist[:]
  except:
    return False

  ncol=10      # number of level indices per row
  lform='%5i'  # output format

  # write to file
  outlines=[]
  while len(lprint) > 0:
    line=''
    for i in range(ncol):
      if len(lprint) == 0: break
      line=line+lform%lprint.pop(0)
    outlines.append(line)
  if len(outlines) == 0:
    print 'No levels'
  else:
    if filename == None:
      for line in outlines: print line
    else:
      fout=open(filename,'w')
      for line in outlines: fout.write(line+'\n')
      fout.close()
  return True

def read_levlist(filename):
  # attempt to open input file
  try:
    fin=open(filename,'r')
  except:
    return False

  # read level list
  out=[]
  try:
    for line in fin:
      tmp=re.findall(r'(\d+)',line)
      for t in tmp: out.append(int(t))
  except:
    return False

  fin.close()
  return out

# *****************************************************************************
# main program
ul=[]
ul.append("usage: %prog [options] filename [filters]")
ul.append("")
ul.append("Return list of level indices which satisfy a collection of")
ul.append("filters.  Selection can be based on configuration, term")
ul.append("(^{2S+1}L), 2J, parity, and energy range (Ry). Output (piped")
ul.append("to a file) can be used in other rtc_ functions to act on a")
ul.append("subset of levels.")
ul.append("")
ul.append("example filters:")
ul.append("  configuration -  c1 c3    (see -i option)")
ul.append("  term          -  2P 4S")
ul.append("  2J            -  3  6")
ul.append("  parity        -  even odd")
ul.append("  energy range  -  [0.0:10.0]  [4,6]")
ul.append("")
ul.append("note: when several filters of the same type are included in")
ul.append("      the list, the effects are OR'ed together.  Filters of")
ul.append("      different types are AND'ed.  An exception exists for")
ul.append("      the energy range; only one range can be specified (the")
ul.append("      last one given).")
usage=""
for u in ul: usage=usage+u+'\n'

parser=OptionParser(usage=usage)
parser.add_option("-i","--info",action="store_true",dest="info",default=False,\
                  help=" print information about adf04 file including "+\
                       "configuration identifiers")
parser.add_option("-e","--energy",action="store_true",dest="energy",\
                  default=False,help="print level energies to screen")
parser.add_option("-l","--levlist",dest="levlist",default=None,\
                  help="input level list (for further filtering)",\
                  metavar="FILE")
parser.add_option("-o","--output",dest="outfile",default=None,\
                  help="output file name",metavar="FILE")
parser.add_option("-c","--coupling",dest="cup",default=None,\
                  help="specify coupling scheme",metavar="CUP")
parser.add_option("-f","--filter",dest="filterfile",default=None,\
                  help="input file containing filter list ",metavar="FILE")
parser.add_option("--eV",action="store_true",dest="useev",default=False,\
                  help="give energies in electron-volts")

(options,args)=parser.parse_args()

if len(args) == 0:
  parser.print_help()
  sys.exit(0)
else:
  filename=args[0]
  filters=args[1:]
  if not os.path.exists(filename):
    print '*** Error - '+str(filename)+' not found ***'
    sys.exit(0)

# use electron-volts instead of Rydbergs in with -e option
useev=False
if options.useev: useev=True

if options.cup not in [None,'IC','LS']:
  print '*** Error - coupling scheme, %s, not supported'%cup
  sys.exit()
if options.cup == 'IC':
  try:
    ion=rap_states_lsj(filename)
    cup='IC'
  except:
    if options.cup == 'IC':
      print '*** Error - no LSJ energy data found ***'
      sys.exit()
elif options.cup == 'LS':
  try:
    ion=rap_states_ls(filename)
    cup='LS'
  except:
    print '*** Error - no LS energy data found ***'
    sys.exit()
else:
  try:
    ion=rlf_loadstates(filename)
    cup=ion.coupling()
  except ErrorRAP, e:
    print e.value
    sys.exit()

levelfile=options.levlist
if levelfile != None:
  if not os.path.exists(levelfile):
    print '*** Error - '+str(levelfile)+' not found ***'
    sys.exit(0)
filterfile=options.filterfile
if filterfile != None:
  if not os.path.exists(filterfile):
    print '*** Error - '+str(filterfile)+' not found ***'
    sys.exit(0)
outfile=options.outfile

if levelfile == None:
  levlist=ion.indices()
else:
  levlist=read_levlist(levelfile)

if filterfile != None:
  fin=open(filterfile,'r')
  for line in fin:
    if line[0] != '#' or line.strip() != '':
      entry=line.strip()
      if entry not in filters: filters.append(entry)
  fin.close()

confs=ion.configurations(levlist)                     # get configuration list
terms=ion.terms(levlist)                              # get term list
(emin,emax)=ion.energy_range(levlist)                 # get energy range
if cup == 'IC': (jmin,jmax)=ion.jrange(levlist)       # get J range

if options.info:
  report_info(ion,levlist,outfile,useev)
elif options.energy:
  levlist=apply_filter(filters,levlist)
  report_energies(outfile,useev)
else:
  levels=apply_filter(filters,levlist)
  write_levlist(levels,outfile)
